package com.example.isoftlangchainai.rag.naive;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;

import java.util.List;

/**
 * @Description: NaivePgVectorRagHelper
 * @Author :chenjun
 */
public class NaivePgVectorRagHelper {
    protected String baseUrl = "http://langchain4j.dev/demo/openai/v1";
    protected String apikey = "demo";
    protected String modelName = "gpt-4o-mini";
    protected EmbeddingStore<TextSegment> embeddingStore;
    public NaivePgVectorRagHelper(EmbeddingStore<TextSegment> embeddingStore) {
        this.embeddingStore = embeddingStore;
    }

    public NaiveRagAssistant createAssistant() {
        System.out.println("1.let's create a chat model:聊天语言模型......");
        ChatModel chatModel = OpenAiChatModel.builder()
                .baseUrl(baseUrl)
                .apiKey(apikey)
                .modelName(modelName)
                .build();

        System.out.println("2.create EmbeddingModel and embed (also known as 'vectorize') these segments......");
        EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

        System.out.println("3.store these embeddings in an embedding store (also known as a 'vector database')......");


        System.out.println("""
               4.创建内容检索器ContentRetriever，用于根据用户查询从向量化存储中检索最相关的文本片段。
                 配置返回相关性最高的结果为2、最低相似度得分为0.7......
                 """);
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(2)
                .minScore(0.7)
                .build();

        Query query = new Query("你喜欢的运动是什么？");
        List<Content> retrieve = contentRetriever.retrieve(query);
        System.out.println("你喜欢的运动是什么:");
        for (Content content : retrieve) {
            System.out.println(content.textSegment().text());
        }

        System.out.println("""
                5.使用聊天记忆（ChatMemory）来支持与大语言模型LLM的多轮对话，使其能够记住之前的交互内容......
                """);
        //按消息数量限制历史记录,如保留最近10条消息
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);

        System.out.println("6.build our AI Service......");
        return AiServices.builder(NaiveRagAssistant.class)
                .chatModel(chatModel)//设置聊天语言模型
                .contentRetriever(contentRetriever)//设置内容检索器
//                .chatMemory(chatMemory)//设置聊天记忆对象
//                .tools()
                .build();
    }

}
